use crate::{
    courier::{codec::*, ConnectionSource, TcpConnectionStream},
    crate_common_types::*,
    utils::{
        auto_abort::with_vec_no_bg_task::AutoAbort, deser_from_decoded,
        maybe_deser_owned::VariantMaybeDeserOwned,
    },
};
use fnv::FnvHashMap;
use futures::{
    ready,
    stream::{AbortHandle, AbortRegistration, Abortable},
    FutureExt, Sink, Stream,
};
use log::{error, info, warn};
use pin_project::pin_project;
use serde::{Deserialize, Serialize};
use std::{
    collections::hash_map::Entry,
    fmt::Debug,
    future::Future,
    io::Error as IoErr,
    net::SocketAddr,
    panic::AssertUnwindSafe,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
};
use tokio::{
    io::{AsyncRead, AsyncWrite}, net::TcpListener, sync::mpsc
};
use tokio_util::codec::Framed;

macro_rules! proj {
    ($it: ident) => {
        $it.as_mut().project()
    };
}

const REQ_ONESHOT_RECV_LIMIT: usize = 100;
const RESP_ONESHOT_SEND_LIMIT: usize = 100;

type RawChannelEnd<C, S, D, I, O> =
    Framed<C, WithFixedSizedPrefix<S, D, RespFramePrefix, ReqFramePrefix, I, O>>;

pub trait ServeRpc: Sync + Send + 'static {
    type OtherEndInfo;
    type Req: Send;
    type Resp: Send;
    fn dispatch_call(
        &self,
        req: Self::Req,
        from: Self::OtherEndInfo,
    ) -> impl Future<Output = Self::Resp> + Send;
}

pub struct Server<S, CS = TcpConnectionStream, Ser = DefaultSer, Deser = DefaultDeser> {
    serve: Arc<S>,
    conn_source: CS,
    ser: Ser,
    deser: Deser,
}

impl<S> Server<S> {
    pub async fn with_default_serder_tcp(
        serve: Arc<S>,
        listen_addr: SocketAddr,
    ) -> Result<Self, IoErr> {
        Ok(Self::with_default_serder(
            serve,
            TcpConnectionStream(TcpListener::bind(listen_addr).await?),
        ))
    }
}

impl<S, CS> Server<S, CS> {
    pub fn with_default_serder(serve: Arc<S>, conn_source: CS) -> Self {
        Self::new(
            serve,
            conn_source,
            DefaultSer::default(),
            DefaultDeser::default(),
        )
    }
}

impl<S, CS, Ser, Deser> Server<S, CS, Ser, Deser> {
    pub fn new(serve: Arc<S>, conn_source: CS, ser: Ser, deser: Deser) -> Self {
        Self {
            serve,
            conn_source,
            ser,
            deser,
        }
    }
}

impl<S, CS, Ser, Deser> Server<S, CS, Ser, Deser>
where
    S: ServeRpc,
    S::Req: Deserialize<'static> + VariantMaybeDeserOwned,
    S::Resp: Serialize + VariantMaybeDeserOwned,
    S::OtherEndInfo: Clone + Send + 'static,
    CS: ConnectionSource,
    CS::OtherEndInfo: Into<S::OtherEndInfo>,
    CS::Conn: Send + 'static,
    Ser: ToWriterSerializer + Clone + Send + 'static,
    Deser: FromSliceDeserializer + Clone + Send + 'static,
    Ser::Err: Debug,
    Deser::Err: Debug,
    S::OtherEndInfo: Debug,
{
    pub async fn serve(self) -> Result<(), CS::Err> {
        let mut auto_abort_set: AutoAbort<512> = AutoAbort::new();
        loop {
            match self.conn_source.get_one().await {
                Ok((conn, info)) => {
                    let raw_end = RawChannelEnd::new(
                        conn,
                        WithFixedSizedPrefix::new(self.ser.clone(), self.deser.clone()),
                    );
                    let (tx, rx) = mpsc::channel(1000);
                    let serving = ServingEnd {
                        other_end_info: info.into(),
                        raw_end,
                        serve: Arc::clone(&self.serve),
                        send_chan_rx: rx,
                        send_chan_tx: tx,
                        ongoing_reqs: Default::default(),
                        auto_abort_set: AutoAbort::new(),
                        deser: self.deser.clone(),
                        receive_from_other_end_done: false,
                    };

                    auto_abort_set.spawn(async move {
                        match AssertUnwindSafe(serving).catch_unwind().await {
                            Err(_e_panic) => {
                                error!("Unexpected panic from a client serving task!")
                            }
                            Ok(Err(e)) => {
                                error!("A client serving task exit with IO error: {e:?}.")
                            }
                            Ok(Ok(_ok)) => {
                                info!("A client serving task is done: the other end closed.")
                            }
                        }
                    });
                }
                Err(e) => break Err(e),
            }
        }
    }
}

/// The serving end of **one** RPC channel.
#[pin_project]
struct ServingEnd<S: ServeRpc, Ser, Deser, C> {
    other_end_info: S::OtherEndInfo,
    #[pin]
    raw_end: RawChannelEnd<C, Ser, Deser, S::Resp, S::Req>,
    serve: Arc<S>,
    send_chan_rx: mpsc::Receiver<(RespFramePrefix, MaybeEmpty<S::Resp>)>,
    ongoing_reqs: FnvHashMap<ClientLocalReqId, AbortHandle>,

    auto_abort_set: AutoAbort<32>,
    send_chan_tx: mpsc::Sender<(RespFramePrefix, MaybeEmpty<S::Resp>)>,
    deser: Deser,

    receive_from_other_end_done: bool,
}

enum PollReceiveDone {
    NothingToReceive,
    OneshotReceiveLimitReached,
}

enum PollSendDone {
    OneshotSendLimitReached,
}

impl<S, Ser, Deser, C> ServingEnd<S, Ser, Deser, C>
where
    S: ServeRpc,
    S::Req: Deserialize<'static> + VariantMaybeDeserOwned,
    S::Resp: Serialize + VariantMaybeDeserOwned,
    S::OtherEndInfo: Clone + Send + 'static,
    Ser: ToWriterSerializer,
    Deser: FromSliceDeserializer + Clone + Send + 'static,
    C: AsyncRead + AsyncWrite,
    Ser::Err: Debug,
    Deser::Err: Debug,
    S::OtherEndInfo: Debug,
{
    fn handle_req(mut self: Pin<&mut Self>, (pre, req): (ReqFramePrefix, Decoded<S::Req>)) {
        match pre.typ {
            ReqType::Call => {
                let Entry::Vacant(entry) = proj!(self).ongoing_reqs.entry(pre.id) else {
                    return;
                };
                let (abort_handle, abort_regist) = AbortHandle::new_pair();
                entry.insert(abort_handle);

                let req_serving = req_overwatch_serve(
                    self.other_end_info.clone(),
                    pre,
                    req,
                    Arc::clone(&self.serve),
                    self.send_chan_tx.clone(),
                    abort_regist,
                    self.deser.clone(),
                );
                proj!(self).auto_abort_set.spawn(req_serving);
            }
            ReqType::Cancel => {
                if let Some(abort_handle) = proj!(self).ongoing_reqs.remove(&pre.id) {
                    abort_handle.abort()
                }
            }
        }
    }
    /// Try to receive:
    /// 1. requests from the underlying data channel
    fn poll_receive(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<PollReceiveDone, IoErr>> {
        if self.receive_from_other_end_done {
            return Poll::Ready(Ok(PollReceiveDone::NothingToReceive))
        }
        for _ in 0..REQ_ONESHOT_RECV_LIMIT {
            match ready!(proj!(self).raw_end.poll_next(cx)) {
                Some(Ok(req)) => self.as_mut().handle_req(req),
                Some(Err(e)) => match e {
                    IoOrDecodeErr::Io(e) => return Poll::Ready(Err(e)),
                    IoOrDecodeErr::PrefixRead(e) => {
                        // TODO: handle prefix read error
                        info!(
                            "Request ? from {:?}: decode prefix error: {e:?}",
                            self.other_end_info
                        )
                    }
                    IoOrDecodeErr::PayloadDeser(pre, e) => {
                        // TODO: handle payload deserialzation error (1)
                        info!(
                            "Request {} from {:?}: deser payload error (1): {e:?}",
                            pre.id, self.other_end_info
                        )
                    }
                    IoOrDecodeErr::ShouldDelegate => {
                        // TODO: handle should delegate error
                        info!("Resp isn't DeserializeOwned, the sender should delegate deser")
                    }
                },
                None => {
                    *proj!(self).receive_from_other_end_done = true;
                    return Poll::Ready(Ok(PollReceiveDone::NothingToReceive))
                },
            }
        }
        Poll::Ready(Ok(PollReceiveDone::OneshotReceiveLimitReached))
    }

    /// Try to send:
    /// 1. responses from the response queue
    fn poll_send(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<PollSendDone, IoErr>> {
        for _ in 0..RESP_ONESHOT_SEND_LIMIT {
            ready!(proj!(self).raw_end.poll_ready(cx)).map_err(IoOrEncodeErr::assert_io_err)?;
            match proj!(self).send_chan_rx.poll_recv(cx) {
                Poll::Ready(None) => unreachable!("Unexpected state: server send chan closed."),
                Poll::Ready(Some((pre, resp))) => {
                    let id = pre.id;
                    if proj!(self).ongoing_reqs.remove(&id).is_none() {
                        continue;
                    }
                    match proj!(self).raw_end.start_send((pre, resp)) {
                        Err(IoOrEncodeErr::Io(e)) => return Poll::Ready(Err(e)),
                        Err(IoOrEncodeErr::PrefixWrite(_)) => {
                            unreachable!("Unexpected prefix error")
                        }
                        Err(IoOrEncodeErr::PayloadSer(e)) => {
                            // TODO: handle payload serialization error
                            warn!(
                                "Request {} from {:?}: ser response payload err: {e:?}",
                                id, self.other_end_info
                            )
                        }
                        Ok(_) => {}
                    };
                }
                Poll::Pending => {
                    ready!(proj!(self).raw_end.poll_flush(cx))
                        .map_err(IoOrEncodeErr::assert_io_err)?;
                    return Poll::Pending;
                }
            }
        }
        Poll::Ready(Ok(PollSendDone::OneshotSendLimitReached))
    }
}

pub enum ServingDone {
    OtherEndClosed,
}

impl<S, Ser, Deser, C> Future for ServingEnd<S, Ser, Deser, C>
where
    S: ServeRpc,
    S::Req: Deserialize<'static> + VariantMaybeDeserOwned,
    S::Resp: Serialize + VariantMaybeDeserOwned,
    S::OtherEndInfo: Clone + Send + 'static,
    Ser: ToWriterSerializer,
    Deser: FromSliceDeserializer + Clone + Send + 'static,
    C: AsyncRead + AsyncWrite,
    Ser::Err: Debug,
    Deser::Err: Debug,
    S::OtherEndInfo: Debug,
{
    type Output = Result<ServingDone, IoErr>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        loop {
            match (
                self.as_mut().poll_receive(cx)?,
                self.as_mut().poll_send(cx)?,
            ) {
                (Poll::Ready(PollReceiveDone::NothingToReceive), poll_send) => {
                    // info!("Client closed connection, #ongoing_reqs = {}", self.ongoing_reqs.len());
                    if self.ongoing_reqs.is_empty() {
                        return Poll::Ready(Ok(ServingDone::OtherEndClosed));
                    }
                    let _ = ready!(poll_send);
                }
                (Poll::Pending, Poll::Pending) => return Poll::Pending,
                _ => {}
            }
        }
    }
}

/// The serving task for **one** request.
async fn req_overwatch_serve<S, Deser>(
    other_end_info: S::OtherEndInfo,
    pre: ReqFramePrefix,
    req: Decoded<S::Req>,
    serve: Arc<S>,
    result_chan: mpsc::Sender<(RespFramePrefix, MaybeEmpty<S::Resp>)>,
    abort_regist: AbortRegistration,
    deser: Deser,
) where
    S: ServeRpc,
    S::Resp: VariantMaybeDeserOwned,
    S::Req: Deserialize<'static> + VariantMaybeDeserOwned,
    Deser: FromSliceDeserializer,
    Deser::Err: Debug,
    S::OtherEndInfo: Debug,
{
    let req_maybe_refed = match deser_from_decoded(req, &deser) {
        Ok(req) => req,
        Err(e) => {
            // TODO: handle payload deserialzation error (2)
            info!(
                "Request {} from {:?}: deser payload error (2): {e:?}",
                pre.id, &other_end_info
            );
            return;
        }
    };
    // SAFETY: we never drop the source before the inner type
    let req_with_src = unsafe { req_maybe_refed.into_inner() };

    // TODO: add task execution time limit, report timeout if the execution
    // takes too long

    let f = Abortable::new(
        AssertUnwindSafe(serve.dispatch_call(req_with_src.0, other_end_info)).catch_unwind(),
        abort_regist,
    );
    let may_send_err = match f.await {
        Ok(Ok(resp)) => {
            result_chan
                .send((
                    RespFramePrefix::new_done(pre.id, !resp.variant_is_deser_owned()),
                    MaybeEmpty::Filled(resp),
                ))
                .await
        }
        Ok(Err(_)) => {
            // the task panicked
            result_chan
                .send((RespFramePrefix::new_panicked(pre.id), MaybeEmpty::Empty))
                .await
        }
        Err(_) => {
            // the task is aborted via the AbortHandle (i.e. cancelled)
            Ok(())
        }
    };

    if may_send_err.is_err() {
        // we never `close` the recv end, and the serving end always hold
        // a recv end, so if we reach here, the serving end must have been
        // droppped
    }
}
